#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PCAP会话拆分工具
基于深度包检测(DPI)技术，按应用层协议拆分PCAP文件中的网络会话
"""

import os
import sys
import json
import time
import argparse
import logging
from datetime import datetime
from pathlib import Path
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List, Tuple, Optional, Set

try:
    from scapy.all import rdpcap, wrpcap, Raw
    from scapy.layers.inet import IP, TCP, UDP
    from scapy.layers.dns import DNS
    try:
        from scapy.layers.http import HTTP, HTTPRequest, HTTPResponse
    except ImportError:
        HTTP = HTTPRequest = HTTPResponse = None
    try:
        from scapy.layers.tls import TLS
    except ImportError:
        TLS = None
except ImportError:
    print("错误: 缺少scapy库，请运行: pip install scapy")
    sys.exit(1)

try:
    from tqdm import tqdm
except ImportError:
    print("错误: 缺少tqdm库，请运行: pip install tqdm")
    sys.exit(1)


class ProtocolDetector:
    """协议检测器 - 基于DPI技术识别应用层协议"""
    
    def __init__(self):
        # HTTP协议特征
        self.http_methods = {b'GET', b'POST', b'PUT', b'DELETE', b'HEAD', b'OPTIONS', b'PATCH'}
        self.http_versions = {b'HTTP/1.0', b'HTTP/1.1', b'HTTP/2.0'}
        
        # FTP协议特征
        self.ftp_commands = {b'USER', b'PASS', b'RETR', b'STOR', b'LIST', b'NLST', b'CWD', b'PWD', b'QUIT'}
        
        # SMTP协议特征
        self.smtp_commands = {b'HELO', b'EHLO', b'MAIL FROM', b'RCPT TO', b'DATA', b'QUIT'}
        
        # POP3协议特征
        self.pop3_commands = {b'USER', b'PASS', b'STAT', b'LIST', b'RETR', b'DELE', b'QUIT'}
        
        # IMAP协议特征
        self.imap_commands = {b'LOGIN', b'SELECT', b'EXAMINE', b'FETCH', b'STORE', b'SEARCH'}
        
        # SSH协议特征
        self.ssh_banner = b'SSH-'
        
        # Telnet协议特征
        self.telnet_commands = {b'\xff\xfb', b'\xff\xfc', b'\xff\xfd', b'\xff\xfe'}  # IAC commands
    
    def detect_protocol(self, packet) -> str:
        """检测数据包的应用层协议"""
        try:
            # 检查是否有应用层数据
            if not packet.haslayer(Raw):
                return self._detect_by_layer(packet)
            
            payload = bytes(packet[Raw].load)
            
            # HTTP检测
            if self._is_http(payload):
                return 'HTTP'
            
            # HTTPS检测 (通过TLS层)
            if TLS and packet.haslayer(TLS):
                return 'HTTPS'
            
            # DNS检测
            if packet.haslayer(DNS):
                return 'DNS'
            
            # FTP检测
            if self._is_ftp(payload):
                return 'FTP'
            
            # SMTP检测
            if self._is_smtp(payload):
                return 'SMTP'
            
            # POP3检测
            if self._is_pop3(payload):
                return 'POP3'
            
            # IMAP检测
            if self._is_imap(payload):
                return 'IMAP'
            
            # SSH检测
            if self._is_ssh(payload):
                return 'SSH'
            
            # Telnet检测
            if self._is_telnet(payload):
                return 'Telnet'
            
            # 基于端口的备用检测
            return self._detect_by_port(packet)
            
        except Exception as e:
            logging.debug(f"协议检测异常: {e}")
            return 'Unknown'
    
    def _detect_by_layer(self, packet) -> str:
        """基于Scapy层检测协议"""
        if HTTP and packet.haslayer(HTTP):
            return 'HTTP'
        elif TLS and packet.haslayer(TLS):
            return 'HTTPS'
        elif packet.haslayer(DNS):
            return 'DNS'
        else:
            return self._detect_by_port(packet)
    
    def _detect_by_port(self, packet) -> str:
        """基于端口的备用检测方法"""
        if packet.haslayer(TCP) or packet.haslayer(UDP):
            sport = packet.sport if hasattr(packet, 'sport') else 0
            dport = packet.dport if hasattr(packet, 'dport') else 0
            
            # 常见端口映射
            port_map = {
                80: 'HTTP', 8080: 'HTTP', 8000: 'HTTP',
                443: 'HTTPS', 8443: 'HTTPS',
                21: 'FTP', 20: 'FTP',
                22: 'SSH',
                23: 'Telnet',
                25: 'SMTP', 587: 'SMTP', 465: 'SMTP',
                110: 'POP3', 995: 'POP3',
                143: 'IMAP', 993: 'IMAP',
                53: 'DNS',
                161: 'SNMP', 162: 'SNMP'
            }
            
            return port_map.get(sport, port_map.get(dport, 'Unknown'))
        
        return 'Unknown'
    
    def _is_http(self, payload: bytes) -> bool:
        """检测HTTP协议"""
        try:
            # 检查HTTP方法
            for method in self.http_methods:
                if payload.startswith(method + b' '):
                    return True
            
            # 检查HTTP响应
            for version in self.http_versions:
                if payload.startswith(version + b' '):
                    return True
            
            # 检查HTTP头部特征
            if b'Host:' in payload or b'User-Agent:' in payload:
                return True
                
        except Exception:
            pass
        return False
    
    def _is_ftp(self, payload: bytes) -> bool:
        """检测FTP协议"""
        try:
            payload_upper = payload.upper()
            for cmd in self.ftp_commands:
                if payload_upper.startswith(cmd + b' ') or payload_upper.startswith(cmd + b'\r\n'):
                    return True
            
            # FTP响应码检测
            if len(payload) >= 3 and payload[:3].isdigit():
                code = int(payload[:3])
                if 100 <= code <= 599:  # FTP响应码范围
                    return True
                    
        except Exception:
            pass
        return False
    
    def _is_smtp(self, payload: bytes) -> bool:
        """检测SMTP协议"""
        try:
            payload_upper = payload.upper()
            for cmd in self.smtp_commands:
                if cmd in payload_upper:
                    return True
                    
        except Exception:
            pass
        return False
    
    def _is_pop3(self, payload: bytes) -> bool:
        """检测POP3协议"""
        try:
            payload_upper = payload.upper()
            for cmd in self.pop3_commands:
                if payload_upper.startswith(cmd + b' '):
                    return True
            
            # POP3响应检测
            if payload.startswith(b'+OK') or payload.startswith(b'-ERR'):
                return True
                
        except Exception:
            pass
        return False
    
    def _is_imap(self, payload: bytes) -> bool:
        """检测IMAP协议"""
        try:
            payload_upper = payload.upper()
            for cmd in self.imap_commands:
                if cmd in payload_upper:
                    return True
                    
        except Exception:
            pass
        return False
    
    def _is_ssh(self, payload: bytes) -> bool:
        """检测SSH协议"""
        return payload.startswith(self.ssh_banner)
    
    def _is_telnet(self, payload: bytes) -> bool:
        """检测Telnet协议"""
        try:
            for cmd in self.telnet_commands:
                if cmd in payload:
                    return True
        except Exception:
            pass
        return False


class SessionManager:
    """会话管理器 - 管理网络会话的生命周期"""
    
    def __init__(self, session_timeout=300):
        self.sessions = {}  # session_key -> session_info
        self.session_timeout = session_timeout
        self.session_counter = 0
    
    def get_session_key(self, packet) -> Optional[str]:
        """生成会话唯一标识"""
        try:
            if packet.haslayer(IP):
                ip_layer = packet[IP]
                src_ip = ip_layer.src
                dst_ip = ip_layer.dst
                
                if packet.haslayer(TCP):
                    tcp_layer = packet[TCP]
                    src_port = tcp_layer.sport
                    dst_port = tcp_layer.dport
                    protocol = 'TCP'
                elif packet.haslayer(UDP):
                    udp_layer = packet[UDP]
                    src_port = udp_layer.sport
                    dst_port = udp_layer.dport
                    protocol = 'UDP'
                else:
                    return None
                
                # 标准化会话key（较小的IP:端口在前）
                if (src_ip, src_port) < (dst_ip, dst_port):
                    return f"{protocol}_{src_ip}_{src_port}_{dst_ip}_{dst_port}"
                else:
                    return f"{protocol}_{dst_ip}_{dst_port}_{src_ip}_{src_port}"
                    
        except Exception as e:
            logging.debug(f"生成会话key失败: {e}")
            return None
    
    def add_packet_to_session(self, packet, protocol: str) -> Optional[str]:
        """将数据包添加到对应会话"""
        session_key = self.get_session_key(packet)
        if not session_key:
            return None
        
        current_time = time.time()
        
        if session_key not in self.sessions:
            # 创建新会话
            self.session_counter += 1
            session_id = f"{self.session_counter:06d}"
            
            self.sessions[session_key] = {
                'session_id': session_id,
                'protocol': protocol,
                'packets': [],
                'start_time': current_time,
                'last_activity': current_time,
                'src_ip': packet[IP].src if packet.haslayer(IP) else '',
                'dst_ip': packet[IP].dst if packet.haslayer(IP) else '',
                'src_port': packet.sport if hasattr(packet, 'sport') else 0,
                'dst_port': packet.dport if hasattr(packet, 'dport') else 0,
                'packet_count': 0,
                'byte_count': 0
            }
        
        # 更新会话信息
        session = self.sessions[session_key]
        session['packets'].append(packet)
        session['last_activity'] = current_time
        session['packet_count'] += 1
        session['byte_count'] += len(packet)
        
        return session_key
    
    def cleanup_expired_sessions(self) -> List[str]:
        """清理过期会话"""
        current_time = time.time()
        expired_sessions = []
        
        for session_key, session in list(self.sessions.items()):
            if current_time - session['last_activity'] > self.session_timeout:
                expired_sessions.append(session_key)
                del self.sessions[session_key]
        
        return expired_sessions
    
    def get_all_sessions(self) -> Dict:
        """获取所有会话"""
        return self.sessions.copy()
    
    def clear_sessions(self):
        """清空所有会话"""
        self.sessions.clear()
        self.session_counter = 0


class OutputManager:
    """输出管理器 - 管理文件输出和目录结构"""
    
    def __init__(self, output_dir: str):
        self.output_dir = Path(output_dir)
        self.protocol_dirs = {}
        self.statistics = defaultdict(int)
        self.session_metadata = []
        
        # 创建输出目录
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        # 创建统计和日志目录
        (self.output_dir / 'statistics').mkdir(exist_ok=True)
        (self.output_dir / 'logs').mkdir(exist_ok=True)
    
    def get_protocol_dir(self, protocol: str) -> Path:
        """获取协议对应的输出目录"""
        if protocol not in self.protocol_dirs:
            protocol_dir = self.output_dir / protocol
            protocol_dir.mkdir(exist_ok=True)
            self.protocol_dirs[protocol] = protocol_dir
        
        return self.protocol_dirs[protocol]
    
    def save_session(self, session_key: str, session_info: Dict):
        """保存会话到文件"""
        try:
            protocol = session_info['protocol']
            protocol_dir = self.get_protocol_dir(protocol)
            
            # 生成文件名
            timestamp = datetime.fromtimestamp(session_info['start_time']).strftime('%Y%m%d_%H%M%S')
            filename = f"{timestamp}_{session_info['src_ip']}_{session_info['src_port']}_to_{session_info['dst_ip']}_{session_info['dst_port']}_{session_info['session_id']}.pcap"
            
            pcap_file = protocol_dir / filename
            
            # 写入PCAP文件
            wrpcap(str(pcap_file), session_info['packets'])
            
            # 生成元数据
            metadata = {
                'session_id': session_info['session_id'],
                'protocol': protocol,
                'src_ip': session_info['src_ip'],
                'src_port': session_info['src_port'],
                'dst_ip': session_info['dst_ip'],
                'dst_port': session_info['dst_port'],
                'start_time': datetime.fromtimestamp(session_info['start_time']).isoformat(),
                'end_time': datetime.fromtimestamp(session_info['last_activity']).isoformat(),
                'duration': session_info['last_activity'] - session_info['start_time'],
                'packet_count': session_info['packet_count'],
                'byte_count': session_info['byte_count'],
                'pcap_file': str(pcap_file.relative_to(self.output_dir))
            }
            
            # 保存元数据文件
            metadata_file = protocol_dir / f"{filename.replace('.pcap', '.json')}"
            with open(metadata_file, 'w', encoding='utf-8') as f:
                json.dump(metadata, f, indent=2, ensure_ascii=False)
            
            # 更新统计信息
            self.statistics[f'{protocol}_sessions'] += 1
            self.statistics[f'{protocol}_packets'] += session_info['packet_count']
            self.statistics[f'{protocol}_bytes'] += session_info['byte_count']
            self.session_metadata.append(metadata)
            
            logging.info(f"保存会话: {filename} ({session_info['packet_count']} 包, {session_info['byte_count']} 字节)")
            
        except Exception as e:
            logging.error(f"保存会话失败 {session_key}: {e}")
    
    def generate_statistics(self):
        """生成统计报告"""
        try:
            # 处理报告
            processing_report = {
                'timestamp': datetime.now().isoformat(),
                'total_sessions': len(self.session_metadata),
                'protocol_distribution': {},
                'total_packets': sum(v for k, v in self.statistics.items() if k.endswith('_packets')),
                'total_bytes': sum(v for k, v in self.statistics.items() if k.endswith('_bytes'))
            }
            
            # 协议分布统计
            protocols = set()
            for key in self.statistics.keys():
                if key.endswith('_sessions'):
                    protocol = key.replace('_sessions', '')
                    protocols.add(protocol)
            
            for protocol in protocols:
                processing_report['protocol_distribution'][protocol] = {
                    'sessions': self.statistics.get(f'{protocol}_sessions', 0),
                    'packets': self.statistics.get(f'{protocol}_packets', 0),
                    'bytes': self.statistics.get(f'{protocol}_bytes', 0)
                }
            
            # 保存报告
            with open(self.output_dir / 'statistics' / 'processing_report.json', 'w', encoding='utf-8') as f:
                json.dump(processing_report, f, indent=2, ensure_ascii=False)
            
            # 保存会话汇总
            with open(self.output_dir / 'statistics' / 'session_summary.json', 'w', encoding='utf-8') as f:
                json.dump(self.session_metadata, f, indent=2, ensure_ascii=False)
            
            logging.info(f"统计报告已生成: {len(self.session_metadata)} 个会话")
            
        except Exception as e:
            logging.error(f"生成统计报告失败: {e}")


class PcapSplitter:
    """PCAP拆分器主类"""
    
    def __init__(self, output_dir: str, session_timeout=300, max_threads=4):
        self.protocol_detector = ProtocolDetector()
        self.session_manager = SessionManager(session_timeout)
        self.output_manager = OutputManager(output_dir)
        self.max_threads = max_threads
        
        # 设置日志
        log_file = Path(output_dir) / 'logs' / 'pcap_splitter.log'
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler()
            ]
        )
    
    def process_pcap_file(self, pcap_file: str, protocols_filter: Optional[Set[str]] = None):
        """处理单个PCAP文件"""
        try:
            logging.info(f"开始处理文件: {pcap_file}")
            
            # 获取文件大小信息
            file_size = Path(pcap_file).stat().st_size
            logging.info(f"文件大小: {file_size / 1024 / 1024:.2f} MB")
            
            # 读取所有数据包但分批处理
            logging.info("正在读取PCAP文件...")
            packets = rdpcap(pcap_file)
            total_packets = len(packets)
            logging.info(f"读取到 {total_packets} 个数据包")
            
            # 分批处理数据包
            batch_size = 5000  # 每批处理5000个包
            packet_count = 0
            
            # 使用进度条显示总体进度
            with tqdm(total=total_packets, desc=f"处理 {Path(pcap_file).name}") as pbar:
                for i in range(0, total_packets, batch_size):
                    batch_packets = packets[i:i+batch_size]
                    batch_num = i // batch_size + 1
                    
                    logging.info(f"处理第 {batch_num} 批，包含 {len(batch_packets)} 个数据包")
                    
                    # 处理当前批次的数据包
                    for packet in batch_packets:
                        try:
                            # 检测协议
                            protocol = self.protocol_detector.detect_protocol(packet)
                            
                            # 协议过滤
                            if protocols_filter and protocol not in protocols_filter:
                                pbar.update(1)
                                continue
                            
                            # 添加到会话
                            session_key = self.session_manager.add_packet_to_session(packet, protocol)
                            packet_count += 1
                            
                            # 定期清理过期会话和保存会话
                            if packet_count % 1000 == 0:
                                expired = self.session_manager.cleanup_expired_sessions()
                                for exp_key in expired:
                                    if exp_key in self.session_manager.sessions:
                                        self.output_manager.save_session(exp_key, self.session_manager.sessions[exp_key])
                            
                            # 内存管理：当会话数量过多时强制保存
                            if len(self.session_manager.sessions) > 3000:
                                logging.info("会话数量过多，执行强制保存...")
                                sessions_to_save = list(self.session_manager.sessions.items())[:1500]
                                for session_key, session_info in sessions_to_save:
                                    self.output_manager.save_session(session_key, session_info)
                                    del self.session_manager.sessions[session_key]
                            
                            pbar.update(1)
                            
                        except Exception as e:
                            logging.debug(f"处理数据包异常: {e}")
                            pbar.update(1)
                            continue
                    
                    # 每批次后清理一次
                    expired = self.session_manager.cleanup_expired_sessions()
                    for exp_key in expired:
                        if exp_key in self.session_manager.sessions:
                            self.output_manager.save_session(exp_key, self.session_manager.sessions[exp_key])
                    
                    logging.info(f"第 {batch_num} 批处理完成，累计处理 {packet_count} 个数据包")
            
            # 保存剩余会话
            logging.info("保存剩余会话...")
            for session_key, session_info in self.session_manager.get_all_sessions().items():
                self.output_manager.save_session(session_key, session_info)
            
            # 清空会话管理器
            self.session_manager.clear_sessions()
            
            logging.info(f"文件处理完成: {pcap_file}，总计处理 {packet_count} 个数据包")
            
        except Exception as e:
            logging.error(f"处理文件失败 {pcap_file}: {e}")
            import traceback
            logging.error(traceback.format_exc())
    
    def process_multiple_files(self, pcap_files: List[str], protocols_filter: Optional[Set[str]] = None):
        """并发处理多个PCAP文件"""
        if len(pcap_files) == 1:
            self.process_pcap_file(pcap_files[0], protocols_filter)
        else:
            with ThreadPoolExecutor(max_workers=min(self.max_threads, len(pcap_files))) as executor:
                futures = []
                for pcap_file in pcap_files:
                    future = executor.submit(self.process_pcap_file, pcap_file, protocols_filter)
                    futures.append(future)
                
                # 等待所有任务完成
                for future in as_completed(futures):
                    try:
                        future.result()
                    except Exception as e:
                        logging.error(f"并发处理异常: {e}")
        
        # 生成最终统计报告
        self.output_manager.generate_statistics()


def find_pcap_files(input_path: str) -> List[str]:
    """查找PCAP文件"""
    pcap_files = []
    input_path = Path(input_path)
    
    if input_path.is_file():
        if input_path.suffix.lower() in ['.pcap', '.pcapng', '.cap']:
            pcap_files.append(str(input_path))
    elif input_path.is_dir():
        for ext in ['*.pcap', '*.pcapng', '*.cap']:
            pcap_files.extend([str(f) for f in input_path.glob(ext)])
    
    return pcap_files


def main():
    parser = argparse.ArgumentParser(description='PCAP会话拆分工具')
    parser.add_argument('-i', '--input', required=True, help='输入PCAP文件或目录')
    parser.add_argument('-o', '--output', required=True, help='输出目录')
    parser.add_argument('-p', '--protocols', help='指定要处理的协议(逗号分隔)')
    parser.add_argument('-t', '--threads', type=int, default=4, help='并发线程数')
    parser.add_argument('--timeout', type=int, default=300, help='会话超时时间(秒)')
    parser.add_argument('-v', '--verbose', action='store_true', help='详细输出')
    
    args = parser.parse_args()
    
    # 设置日志级别
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    # 查找PCAP文件
    pcap_files = find_pcap_files(args.input)
    if not pcap_files:
        print(f"错误: 在 {args.input} 中未找到PCAP文件")
        sys.exit(1)
    
    print(f"找到 {len(pcap_files)} 个PCAP文件")
    for f in pcap_files:
        print(f"  - {f}")
    
    # 解析协议过滤器
    protocols_filter = None
    if args.protocols:
        protocols_filter = set(p.strip().upper() for p in args.protocols.split(','))
        print(f"协议过滤器: {protocols_filter}")
    
    # 创建拆分器并处理文件
    splitter = PcapSplitter(
        output_dir=args.output,
        session_timeout=args.timeout,
        max_threads=args.threads
    )
    
    start_time = time.time()
    splitter.process_multiple_files(pcap_files, protocols_filter)
    end_time = time.time()
    
    print(f"\n处理完成! 耗时: {end_time - start_time:.2f} 秒")
    print(f"输出目录: {args.output}")


if __name__ == '__main__':
    main()